package com.github.tommas;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.Opcodes;
import com.github.tommas.decompiler.Decompiler;
import com.github.tommas.decompiler.FernFlowerDecompiler;

import java.io.File;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.nio.file.FileVisitResult;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.SimpleFileVisitor;
import java.nio.file.attribute.BasicFileAttributes;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class DynamicClassGenerator {
    private final static String DynamicProxyClassNameRegex = "([a-zA-Z_][a-zA-Z0-9_]*)(/)?\\$Proxy\\d+";
    private final static String SpringFrameworkGenerateClassNameRegex = "\\$\\$(.+)CGLIB\\$\\$";

    Decompiler decompiler;
    List<Pattern> classNamePatterns;

    List<String> targetPackages;

    public DynamicClassGenerator(GeneratorOptions options) {
        File outDir = new File(options.outDir);
        if (outDir.exists()) {
            try {
                cleanOutDirectory(outDir.toPath());
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
        decompiler = new FernFlowerDecompiler(options.outDir);

        List<Pattern> patterns = new ArrayList<>();
        patterns.add(Pattern.compile(DynamicProxyClassNameRegex));
        patterns.add(Pattern.compile(SpringFrameworkGenerateClassNameRegex));
        classNamePatterns = patterns;

        if (options.packages != null && !options.packages.isEmpty()) {
            targetPackages = new ArrayList<>(options.packages.size());
            for (String p : options.packages) {
                targetPackages.add(p.replace(".", "/"));
            }
        } else {
            targetPackages = Collections.emptyList();
        }
    }

    static void cleanOutDirectory(Path outPath) throws IOException {
        Files.walkFileTree(outPath, new SimpleFileVisitor<Path>() {
            @Override
            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
                Files.delete(file);
                return FileVisitResult.CONTINUE;
            }

            @Override
            public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException {
                if (exc == null) {
                    Files.delete(dir);
                    return FileVisitResult.CONTINUE;
                } else {
                    throw exc;
                }
            }
        });
    }

    void generate(String className, byte[] classfileBuffer) {
        for (Pattern p : classNamePatterns) {
            Matcher m = p.matcher(className);
            if (m.find()) {
                if (!targetPackages.isEmpty()) {
                    ClassFinder cf = new ClassFinder(targetPackages);
                    ClassReader cr = new ClassReader(classfileBuffer);
                    cr.accept(cf, ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
                    if (!cf.result) {
                        return;
                    }
                }

                decompiler.decompile(className, classfileBuffer);
                return;
            }
        }
    }

    static class ClassFinder extends ClassVisitor {
        List<String> targetPackages;

        boolean result;
        public ClassFinder(List<String> targetPackages) {
            super(Opcodes.ASM4);
            this.targetPackages = targetPackages;
        }

        @Override
        public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
            for (String targetPackage : targetPackages) {
                if (name.startsWith(targetPackage)) {
                    result = true;
                    return;
                }

                if (superName.startsWith(targetPackage)) {
                    result = true;
                    return;
                }

                for (String in : interfaces) {
                    if (in.startsWith(targetPackage)) {
                        result = true;
                        return;
                    }
                }
            }
        }
    }

    public static class Agent {
        public static void premain(String args, Instrumentation instrumentation) {
            final DynamicClassGenerator generator = new DynamicClassGenerator(GeneratorOptions.parse(args));
            instrumentation.addTransformer(new ClassFileTransformer() {
                @Override
                public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                    generator.generate(className, classfileBuffer);
                    return null;
                }
            }, true);
        }
    }
}
